import time

from instruments.instrument import Instrument


class SignalAnalyzer(Instrument):
    def __init__(self, address, parent=None):
        super().__init__(address, parent)


class KeysightSignalAnalyzer(SignalAnalyzer):
    def __init__(self, address, parent=None):
        super().__init__(address, parent)

    def set_frequency(self, freq):
        self._inst.write(':SENSe:FREQuency:CENTer {0}'.format(freq))

    def set_span(self, span):
        self._inst.write_cmd(':SENSe:FREQuency:SPAN {0}'.format(span))

    def set_rbw(self, rbw):
        self._inst.write_cmd(':SENSe:BANDwidth:RESolution:AUTO OFF')
        self._inst.write_cmd(':SENSe:BANDwidth:RESolution {0}'.format(rbw))

    def set_vbw(self, vbw):
        self._inst.write_cmd(':SENSe:BANDwidth:VIDeo:AUTO OFF')
        self._inst.write_cmd(':SENSe:BANDwidth:VIDeo {0}'.format(vbw))

    def set_ref_level(self, level):
        self._inst.write('DISP:WIND:TRAC:Y:RLEV {0} dBm'.format(level))

    def set_att(self, att):
        if str(att) == 'AUTO':
            self._auto_att('ON')
        else:
            self._inst.write_cmd(':SENSe:POWer:ATTenuation {0}'.format(att))

    def set_sweep_time(self, sweep_time):
        if sweep_time == 'AUTO':
            self.set_sweep_time_auto('ON')
            return
        self._inst.write_cmd(':SENSe:SWEep:TIME {0}'.format(sweep_time))

    def set_sweep_time_auto(self, state):
        self._inst.write_cmd(':SENSe:SWEep:TIME:AUTO {0}'.format(state))

    def set_trace_type(self, trace_type):
        self._inst.write_cmd(f'TRAC:TYPE {trace_type}')

    def set_detector(self, detector):
        self._inst.write_cmd(':SENSe:DETector:FUNCtion {0}'.format(detector))

    def restart(self):
        self._inst.write_cmd(':INIT:REST')

    def set_mark_frequency_or_time(self, mark_num, frequency_or_time, marker_type='POS'):
        scpi = f':CALCulate:MARKer{mark_num}:X {frequency_or_time}'
        self._inst.write_cmd(scpi)

    def _auto_att(self, state):
        scpi = f"POW:ATT:AUTO {state}"
        self._inst.write_cmd(scpi)

    def preset(self):
        self._inst.query_cmd(':SYSTem:PRESet;*OPC?')

    def get_power(self, mark_num=1):
        # 标记最高点
        self.xapp_connectivity_set_sa_peak_status(mark_num, 1)
        self._inst.write(':CALC:MARK{0}:MAX'.format(str(mark_num)))
        time.sleep(0.5)  # 解决因速度过快导致标记出现异常。
        # 读取功率值
        power = self.get_mark_power(mark_num)
        return power

    def xapp_connectivity_set_sa_peak_status(self, mark_num, status):
        status_list = ['OFF', 'ON']
        self._inst.write_cmd(':CALC:MARK{0}:STAT {1}'.format(mark_num, status_list[status]))

    def get_mark_power(self, mark_num):
        """
            >> CALC:MARK1:Y?
            << -8.717E+00 (即-8.72 dBm)
        """
        scpi = f':CALC:MARK{mark_num}:Y?'
        power_str = self._inst.query_cmd(scpi)
        # power = float(power_str)
        power =  round(float(power_str), 2)
        return power


class RSSignalAnalyzer(SignalAnalyzer):
    def __init__(self, address, parent=None):
        super().__init__(address, parent)

    def set_frequency(self, freq):
        self._inst.write('FREQ:CENT ' + freq)

    def set_span(self, span):
        self._inst.write(':FREQ:SPAN ' + str(span))

    def set_rbw(self, rbw):
        self._inst.write_cmd('BAND:RES {}'.format(rbw))

    def set_vbw(self, vbw):
        self._inst.write_cmd('BAND:VID {}'.format(vbw))

    def set_ref_level(self, level):
        self._inst.write_cmd('DISP:TRAC:Y:RLEV {}'.format(level))

    def set_att(self, att):
        if str(att) == 'AUTO':
            self._auto_att('ON')
        else:
            self._inst.write_cmd('INP:ATT {}'.format(att))

    def set_sweep_time(self, sweep_time):
        if sweep_time == 'AUTO':
            self.set_sweep_time_auto('ON')
            return
        self._inst.write_cmd('SWE:TIME {}'.format(sweep_time))

    def set_trace_type(self, trace_type):
        self._inst.write_cmd(f'DISP:TRAC1:MODE {trace_type}')

    def set_detector(self, detector):
        self._inst.write_cmd('DET {}'.format(detector))

    def restart(self):
        scpi = 'INIT'
        self._inst.write(scpi)

    def set_mark_frequency_or_time(self, mark_num, frequency_or_time, marker_type='POS'):
        if marker_type == 'DELTa':
            scpi = f":CALC:DELT{mark_num}:X {frequency_or_time}"
        else:
            scpi = f':CALCulate:MARKer{mark_num}:X {frequency_or_time}'
        self._inst.write_cmd(scpi)

    def set_sweep_time_auto(self, state):
        self._inst.write_cmd(f':SENSe:SWEep:TIME:AUTO {state}')

    def _auto_att(self, state):
        scpi = f"INP:ATT:AUTO {state}"
        self._inst.write_cmd(scpi)

    def preset(self):
        self._inst.write(':SYSTem:PRESet')


    def xapp_connectivity_get_sa_peak(self, mark_num=1):
        time.sleep(0.5)
        self._inst.write('CALC:MARK{0}:MAX'.format(str(mark_num)))

    def get_power(self, mark_num=1):
        self.xapp_connectivity_get_sa_peak(mark_num)
        peak_power = self.get_mark_power(mark_num)
        return peak_power

    def get_mark_power(self, mark_num):
        """
            >> CALC:MARK1:Y?
            << -8.717E+00 (即-8.72 dBm)
        """
        scpi = f':CALC:MARK{mark_num}:Y?'
        power_str = self._inst.query_cmd(scpi)
        power = round(float(power_str), 2)  # 邦测改成小数点后两位
        return power